import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

def predict(messages, model, tokenizer):
    device = "cuda"

    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    model_inputs = tokenizer([text], return_tensors="pt").to(device)

    generated_ids = model.generate(model_inputs.input_ids, max_new_tokens=512)
    generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]
    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

    return response


# 加载原下载路径的tokenizer和model
tokenizer = AutoTokenizer.from_pretrained("./qwen/Qwen2-1___5B-Instruct/", use_fast=False, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("./qwen/Qwen2-1___5B-Instruct/", device_map="auto", torch_dtype=torch.bfloat16)

# 加载训练好的Lora模型，将下面的checkpointXXX替换为实际的checkpoint文件名名称
model = PeftModel.from_pretrained(model, model_id="./output-2/Qwen1.5/checkpoint-500")

test_texts = {
    'instruction': "你是一个文本分类领域的专家，你会接收到一段文本和几个潜在的分类选项，请输出文本内容的正确类型",
    'input': "文本: 科技期刊竞技体育在两种经济体制下运行特点分析鲍明晓(北京体育师范学院，北京 100088)摘要 研究主要探讨竞技体育在计划经济体制下和市场经济体制下的运行特点，认为竞技体育在两种经济体制下运行各有优势与不足，建立一个与社会主义市场经济体制相适应的、符合现代竞技运动发展规律的新型竞技体育管理体制和运行机制，必须批判地继承和借鉴竞技体育在两种经济体制下运行的优势与长处。关键词 竞技体育 计划经济体制市场经济体制 运行特点Discussion on the Characteristics of Competitive Sport UnderBoth PlannedEconomic and Market Economic SystemsBao Mingxiao 1998, 18 (1)：6(Beijing Teachers College of Physical Education, Beijing, China 100088)Abstract Development characteristics of  competitive sportsunder both planned economic and market economic  systems  werediscussed and analyzed with the  conclusion  that both systems haveits positive and negative effects  on  the  development of competitivesports. The establishment of a new management system andoperating mechanism for competitive sport, which both suits to thesocialist market economic system  and  meets the  rules ofmodem  sports development, thus must critically inherit and assimilate  the strong points and advantages of both the systems.Key words  sport, administration, economics, comparativestudy, management system, fund raising1 竞技体育在计划经济体制下的运行竞技体育在不同经济体制下有不同的运行方式，这是社会的物质生产方式决定精神生产方式的必然结果。竞技体育在计划经济体制下运行主要有以下几方面的特点：1.1 政治利益驱动采用计划经济体制的国家多是新兴社会主义国家，这些国家在建国伊始就面临十分严酷的生存环境。类型选型:['Energy', 'Agriculture', 'Law', 'Economy', 'Philosophy', 'Military', 'Computer', 'Transport', 'Communication', 'Space', 'Electronics', 'History', 'Sports', 'Medical', 'Enviornment', 'Art', 'Literature', 'Politics']。"
}

instruction = test_texts['instruction']
input_value = test_texts['input']

messages = [
    {"role": "system", "content": f"{instruction}"},
    {"role": "user", "content": f"{input_value}"}
]

response = predict(messages, model, tokenizer)
print(response)
